import cv2
import numpy as np
from scipy.optimize import minimize

def preprocess_image(gray):
    # 将灰度图并归一化
    # gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return (gray - np.mean(gray)) / np.std(gray)

def compute_ncc(img1, img2):
    # 计算归一化互相关NCC
    return np.sum(img1 * img2) / np.sqrt(np.sum(img1**2) * np.sum(img2**2))

def objective_function(params, ref_image, target_image):
    # 目标函数：计算 NCC 相似度"""
    tx, ty = params  # 平移参数
    height, width = ref_image.shape
    # 平移图像
    translated_img = np.zeros_like(ref_image)
    x1 = max(0, int(tx))
    y1 = max(0, int(ty))
    x2 = min(width, width + int(tx))
    y2 = min(height, height + int(ty))
    translated_img[y1:y2, x1:x2] = target_image[y1 - int(ty):y2 - int(ty), x1 - int(tx):x2 - int(tx)]
    # 计算 NCC
    return -compute_ncc(ref_image, translated_img)  # 最大化 NCC

# 基于 NCC 的图像对齐
def align_images(ref_image, target_image):
    # 预处理图像
    ref_gray = preprocess_image(ref_image)
    target_gray = preprocess_image(target_image)

    # 优化平移参数
    initial_guess = [0, 0]  # 初始平移参数
    result = minimize(objective_function, initial_guess, args=(ref_gray, target_gray), method='Powell')
    tx, ty = result.x

    # 应用平移变换
    M = np.float32([[1, 0, tx], [0, 1, ty]])
    aligned_image = cv2.warpAffine(target_image, M, (ref_image.shape[1], ref_image.shape[0]))

    return aligned_image, tx, ty

if __name__ == "__main__":
    # 读取图像为灰度图
    ref_image = cv2.imread('pic_input/src2.jpg', cv2.IMREAD_GRAYSCALE)
    target_image = cv2.imread('pic_input/draw2.jpg', cv2.IMREAD_GRAYSCALE)

    # 检查图像是否成功读取
    if ref_image is None:
        print("Error: Failed to load reference image. Check the file path.")
    if target_image is None:
        print("Error: Failed to load target image. Check the file path.")


    # 对齐图像
    aligned_image, tx, ty = align_images(ref_image, target_image)

    # 保存结果
    cv2.imwrite('pic_input/align_image2.jpg', aligned_image)

    # 显示结果
    # cv2.imshow("Reference Image", ref_image)
    # cv2.imshow("Target Image", target_image)
    # cv2.imshow("Aligned Image", aligned_image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    print(f"Translation Parameters: tx={tx:.2f}, ty={ty:.2f}")